from core.variable import Variable
from utils.common import pair, to_variable
from utils.functions_collect import im2col, linear
from utils.functions_conv import get_conv_outsize

import numpy as np

def conv2d_simple(x, W, b=None, stride=1, pad=0):
    x, W = to_variable(x), to_variable(W)

    Weight = W
    N, C, H, W = x.shape
    OC, C, KH, KW = Weight.shape
    SH, SW = pair(stride)
    PH, PW = pair(pad)
    OH = get_conv_outsize(H, KH, SH, PH)
    OW = get_conv_outsize(W, KW, SW, PW)

    col = im2col(x, (KH, KW), stride, pad, to_matrix=True)
    Weight = Weight.reshape(OC, -1).transpose()
    t = linear(col, Weight, b)
    y = t.reshape(N, OH, OW, OC).transpose(0, 3, 1, 2)
    return y

# im2col
x1 = np.random.rand(1, 3, 7, 7)
col1 = im2col(x1, kernel_size=5, stride=1, pad=0, to_matrix=True)
print(col1.shape)  # (9, 75)

x2 = np.random.rand(10, 3, 7, 7)  # 10个数据
kernel_size = (5, 5)
stride = (1, 1)
pad = (0, 0)
col2 = im2col(x2, kernel_size, stride, pad, to_matrix=True)
print(col2.shape)  # (90, 75)


# conv2d
N, C, H, W = 1, 5, 15, 15
OC, (KH, KW) = 8, (3, 3)
x = Variable(np.random.randn(N, C, H, W))
W = np.random.randn(OC, C, KH, KW)
y = conv2d_simple(x, W, b=None, stride=1, pad=1)
y.backward()
print(y.shape)  # (1, 8, 15, 15)
print(x.grad.shape)  # (1, 5, 15, 15)
